import pandas as pd
import numpy as np
from openpyxl.utils import get_column_letter
import time

BASE_START_TIME = 82500000
BASE_START_TIME_HEADER = (8 * 3600 + 25 * 60) * 1000


# 装饰器, 用于计算函数执行时间
def timing_closure(metric_name):
    """性能监控闭包工厂"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            start = time.time()
            result = func(*args, **kwargs)
            elapsed = time.time() - start
            print(f"[{metric_name}] 耗时: {elapsed:.3f}秒")
            return result
        return wrapper
    return decorator


# @timing_closure("generate_time_ranges")
def generate_time_ranges(interval):
    """生成固定起始时间的时间区间"""
    interval_ms = interval * 60_000
    start = BASE_START_TIME
    end = start + interval_ms - 1  # 确保区间覆盖59秒999毫秒
    ranges = []
    
    # 添加交易日结束时间作为终止条件（假设15:00:00收盘）
    TRADING_DAY_END = 150000000  # 请根据实际收盘时间调整该值
    
    while start <= TRADING_DAY_END:
        current_end = min(end, TRADING_DAY_END)  # 防止生成超过收盘时间的区间
        ranges.append((start, current_end))
        start += interval_ms
        end += interval_ms
    return ranges


# @timing_closure("format_time_range")
def format_time_range(start_ms, end_ms):
    """将毫秒时间戳转换为时间区间字符串"""
    start = pd.to_datetime(start_ms, unit='ms').strftime('%H:%M:%S.%f')[:-3]
    end = pd.to_datetime(end_ms, unit='ms').strftime('%H:%M:%S.%f')[:-3]
    return f"{start} - {end}"


@timing_closure("calculate_kline")
def calculate_kline(data, interval, ranges):
    """基于固定时间区间的K线计算， 当前只会计算1分钟的K线，因为5分钟，10分钟，30分钟的都可以基于1分钟的进行计算，避免了重复计算"""
    data.loc[:, 'time_bin'] = ((data['nTime'] - BASE_START_TIME) // (interval * 60_000)).clip(lower=0)
    
    kline = data.groupby('time_bin').agg(
        open=('nMatch', 'first'),
        high=('nMatch', 'max'),
        low=('nMatch', 'min'),
        close=('nMatch', 'last')
    ).reset_index()
    
    # 映射时间区间
    kline['time_range'] = kline['time_bin'].apply(
        lambda x: format_time_range(*ranges[x]) if x < len(ranges) else None
    )
    return kline[['time_range', 'open', 'high', 'low', 'close']]


# @timing_closure("derive_kline")
def derive_kline(base_kline, base_interval, target_interval):
    """从基础K线推导更高周期， 当前基于1分钟K线进行推导"""
    factor = target_interval // base_interval
    derived = base_kline.groupby(np.arange(len(base_kline)) // factor).agg({
        'open': 'first',
        'high': 'max',
        'low': 'min',
        'close': 'last',
        'time_range': lambda x: f"{x.iloc[0].split(' - ')[0]} - {x.iloc[-1].split(' - ')[1]}"
    })
    return derived


def generate_time_ranges_for_header(interval):
    """生成固定起始时间的时间区间（修复时区基准）"""
    interval_ms = interval * 60_000
    start = BASE_START_TIME_HEADER
    end = start + interval_ms - 1
    ranges = []
    
    # 修正为完整的4小时交易时间
    TRADING_DAY_END = BASE_START_TIME_HEADER + 4 * 3600 * 1000  # 8:25 - 12:25
    
    while start <= TRADING_DAY_END:
        current_end = min(end, TRADING_DAY_END)
        ranges.append((start, current_end))
        start += interval_ms
        end += interval_ms
    return ranges


def format_time_range_for_header(start_ms, end_ms):
    """直接使用本地时间转换（移除时区转换）"""
    start = pd.to_datetime(start_ms, unit='ms').strftime('%H:%M:%S.%f')[:-3]
    end = pd.to_datetime(end_ms, unit='ms').strftime('%H:%M:%S.%f')[:-3]
    return f"{start} - {end}"


def generate_global_header(intervals, output_path, writer):
    global_ranges = generate_time_ranges(1)  # 只生成一次基础时间区间
    
    # 预先创建所有sheet并写入标题
    for sheet_name, interval in intervals.items():
        # 预处理所有时间区间标题
        time_headers = [format_time_range_for_header(*r) for r in generate_time_ranges_for_header(interval)]
        header_df = pd.DataFrame(columns=['股票编号', 'Kline'] + time_headers)
        header_df.to_excel(writer, sheet_name=sheet_name, index=False)
    return global_ranges


def create_excel_with_sheets(csv_path, output_path, specified_codes=None):
    start_time = time.time()
    # 读取CSV数据
    df = pd.read_csv(csv_path, skiprows=1, names=['nMatch', 'nTime', 'nTradingDay', 'szWindCode'])
    
    # 生成所有时间区间
    intervals = {
        '1min Kline': 1,
        '5min Kline': 5,
        '10min Kline': 10,
        '30min Kline': 30
    }

    # 打印 unique_codes 的大小
    unique_codes_num = len(df['szWindCode'].unique())
    print(f"unique_codes 的大小为: {unique_codes_num}")

    grouped = df.groupby('szWindCode')

    with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
        global_ranges = generate_global_header(intervals, output_path, writer)
        for code in df['szWindCode'].unique():
            start_time_code = time.time()  # 记录代码开始时间
            if specified_codes and code not in specified_codes:
                continue

            try:
                stock_data = grouped.get_group(code).copy()
            except KeyError:
                continue
            
            # 计算基础1分钟K线
            base_kline = calculate_kline(stock_data, 1, global_ranges)
            
            start_time_write = time.time()  # 记录写入开始时间
            # 生成各周期K线
            for sheet_name, interval in intervals.items():
                if interval == 1:
                    kline = base_kline
                else:
                    # 除了一分钟K线，其他的可以通过推到方式进行计算
                    kline = derive_kline(base_kline, 1, interval)
                
                # 格式化输出
                time_ranges = kline['time_range'].tolist()
                data = {
                    '股票编号': [code] * 4,
                    'Kline': ['Open', 'High', 'Low', 'Close']
                }
                
                # 为每个时间区间添加四行数据
                for i, tr in enumerate(time_ranges):
                    data[tr] = [
                        kline.iloc[i]['open'],
                        kline.iloc[i]['high'],
                        kline.iloc[i]['low'],
                        kline.iloc[i]['close']
                    ]

                output_df = pd.DataFrame(data)
                output_df.to_excel(
                    writer,
                    sheet_name=sheet_name,
                    index=False,
                    header=False,
                    startrow=writer.sheets[sheet_name].max_row
                )
            end_time_write = time.time()  # 记录写入结束时间
            execution_time_write = end_time_write - start_time_write  # 计算写入执行时间
            print(f"处理 {code} 的写入时间: {execution_time_write:.3f} 秒")
            end_time_code = time.time()  # 记录代码结束时间
            execution_time_code = end_time_code - start_time_code  # 计算代码执行时间
            print(f"处理 {code} 的总时间: {execution_time_code:.3f} 秒")
    end_time = time.time()  # 记录结束时间
    execution_time = end_time - start_time  # 计算执行时间
    print(f"处理所有数据的总时间: {execution_time:.3f} 秒")


if __name__ == "__main__":
    # specified_codes 用来做简单验证，如果指定code，则会只处理指定的code，否则会处理所有的code
    specified_codes = ['104535.SZ', '002412.SZ', '002424.SZ', '000001.SZ', '123110.SZ']
    # 第一个参数是输入的csv文件路径，第二个参数是输出的excel文件路径，第三个参数是指定的code列表，默认为None
    create_excel_with_sheets('md_20221110_appr_500.csv', 'result_2.xlsx')